[TRTLLM-14511][feat] BREAKING: refactor per-request perf metrics for multi-process trtllm-server#16668
[TRTLLM-14511][feat] BREAKING: refactor per-request perf metrics for multi-process trtllm-server#16668reasonsolo wants to merge 10 commits into
Conversation
c0b42bf to
598991a
Compare
reasonsolo
changed the title
reasonsolo
added
the
api-compatible
402269a to
5d92a7d
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughReplaces the HTTP ChangesServer-side JSONL performance metrics pipeline
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant OpenAIHttpClient
participant OpenAIDisaggServer
participant PerfMetricsMiddleware
participant PerfMetricsJsonlWriter
OpenAIHttpClient->>OpenAIDisaggServer: send metrics-enabled request
OpenAIDisaggServer->>PerfMetricsMiddleware: collect completed request records
PerfMetricsMiddleware->>OpenAIHttpClient: return headers or SSE metrics event
PerfMetricsMiddleware->>PerfMetricsJsonlWriter: submit JSONL record
PerfMetricsJsonlWriter->>PerfMetricsJsonlWriter: write per-process JSONL
Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
/bot run --disable-fail-fast |
|
PR_Github #61165 [ run ] triggered by Bot. Commit: |
|
PR_Github #61165 [ run ] completed with state
|
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
d9ff361 to
4f9117c
Compare
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #61270 [ run ] triggered by Bot. Commit: |
|
PR_Github #61270 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61508 [ run ] triggered by Bot. Commit: |
|
PR_Github #61508 [ run ] completed with state
|
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
reasonsolo
added
api-breaking
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/llmapi/apps/test_request_metrics.py (1)
79-89: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winRegister the changed
tests/unittest/llmapi/apps/test_request_metrics.pycases in the integration test lists, or explain why they’re excluded. No matching entries were found intests/integration/test_lists/test-db/*ortests/integration/test_lists/qa/*, and there’s no CBTS touchmap/report to establish coverage. Changed tests: addedtest_time_breakdown_parser_accepts_header_derived_disagg_record,test_jsonl_writer_drops_only_malformed_record; modifiedtest_metrics_headers_use_metric_list_syntax,test_combine_disagg_metrics_is_request_local,test_middleware_controls_public_headers,test_middleware_limits_non_streaming_metrics_headers,test_stream_metrics_follow_done,test_disconnect_after_done_is_ignored,test_file_middleware_intercepts_detail_headers. Coverage verdict: needs follow-up.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/llmapi/apps/test_request_metrics.py` around lines 79 - 89, Register all listed changed tests from test_request_metrics.py in the appropriate integration test lists under test-db or qa, or add a CBTS touchmap/report documenting their exclusion and coverage rationale. Ensure every added or modified test has an explicit coverage entry before completing the change.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tests/unittest/llmapi/apps/test_request_metrics.py`:
- Around line 79-89: Register all listed changed tests from
test_request_metrics.py in the appropriate integration test lists under test-db
or qa, or add a CBTS touchmap/report documenting their exclusion and coverage
rationale. Ensure every added or modified test has an explicit coverage entry
before completing the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6b147191-466e-4fe7-bb1e-092fcd6ed087
📒 Files selected for processing (4)
tensorrt_llm/serve/openai_disagg_server.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/perf_metrics.pytests/unittest/llmapi/apps/test_request_metrics.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/serve/openai_disagg_server.py
- tensorrt_llm/serve/perf_metrics.py
- tensorrt_llm/serve/openai_server.py
Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/disaggregated/test_openai_disagg_service.py`:
- Line 407: Replace assert_not_awaited() with assert_not_called() for
service._gen_client.send_request so the test verifies the method is never
invoked, not merely never awaited.
- Around line 385-407: Update test_context_only_response_finishes_hooks to use
assert_not_called() on service._gen_client.send_request instead of
assert_not_awaited(), ensuring any invocation of the generation client is
detected regardless of whether it was awaited.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9eb4690b-6f6d-4620-8138-df3798cd5e5b
📒 Files selected for processing (2)
tensorrt_llm/serve/openai_disagg_service.pytests/unittest/disaggregated/test_openai_disagg_service.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/serve/openai_disagg_service.py
| @pytest.mark.asyncio | ||
| @pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) | ||
| async def test_context_only_response_finishes_hooks(stream): | ||
| service = _make_service("context_first") | ||
| service._ctx_client = AsyncMock() | ||
| service._gen_client = AsyncMock() | ||
| service._coordinator.get_disagg_request_id = AsyncMock(return_value=42) | ||
| service._check_conditional_disagg = AsyncMock(return_value=(None, True)) | ||
| service._check_gen_only_disagg = AsyncMock(return_value=False) | ||
| service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) | ||
| ctx_response = _make_completion_response("done", finish_reason="stop", context_only=True) | ||
| service._ctx_client.send_request = AsyncMock(return_value=ctx_response) | ||
| hooks = mock.Mock() | ||
| request = CompletionRequest(model="test-model", prompt="hello", stream=stream) | ||
|
|
||
| result = await service._send_disagg_request(request, hooks) | ||
| if stream: | ||
| assert [chunk async for chunk in result] == [b"data: [DONE]\n\n"] | ||
| else: | ||
| assert result is ctx_response | ||
|
|
||
| hooks.on_resp_done.assert_called_once_with("", request, ctx_response) | ||
| service._gen_client.send_request.assert_not_awaited() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
for dir in tests/integration/test_lists/test-db tests/integration/test_lists/qa; do
if [ -d "$dir" ]; then
echo "== $dir =="
rg -n "test_context_only_response_finishes_hooks|test_openai_disagg_service" "$dir" || true
else
echo "Missing: $dir"
fi
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 346
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the test and the implementation path it exercises.
python3 - <<'PY'
from pathlib import Path
paths = [
"tests/unittest/disaggregated/test_openai_disagg_service.py",
"tensorrt_llm/serving/disaggregated/openai_service.py",
]
for p in paths:
path = Path(p)
print(f"\n== {p} ==")
if not path.exists():
print("missing")
continue
lines = path.read_text().splitlines()
for i, line in enumerate(lines, 1):
if "test_context_only_response_finishes_hooks" in line or "_send_disagg_request" in line or "send_request" in line:
start = max(1, i - 20)
end = min(len(lines), i + 40)
print(f"-- context around line {i} --")
for j in range(start, end + 1):
print(f"{j:4d}: {lines[j-1]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Probe unittest.mock.AsyncMock assertion semantics directly.
python3 - <<'PY'
from unittest.mock import AsyncMock
m = AsyncMock()
print("initial:", m.called, m.await_count)
try:
m.assert_not_awaited()
print("assert_not_awaited: pass before call")
except Exception as e:
print("assert_not_awaited: fail before call", type(e).__name__, e)
_ = m("x")
print("after call:", m.called, m.await_count)
try:
m.assert_not_awaited()
print("assert_not_awaited: pass after call")
except Exception as e:
print("assert_not_awaited: fail after call", type(e).__name__, e)
# create a fresh mock and await it
async def main():
n = AsyncMock()
await n("y")
print("after await:", n.called, n.await_count)
try:
n.assert_not_awaited()
print("assert_not_awaited: pass after await")
except Exception as e:
print("assert_not_awaited: fail after await", type(e).__name__, e)
import asyncio
asyncio.run(main())
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 482
Use assert_not_called() for the generation client
assert_not_awaited() still passes if send_request() is invoked but never awaited, so this test can miss an unwanted generation request.
Test coverage summary: test_context_only_response_finishes_hooks is added and listed in tests/integration/test_lists/test-db/l0_a10.yml; no matching qa/ entry. Coverage verdict: sufficient.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/disaggregated/test_openai_disagg_service.py` around lines 385
- 407, Update test_context_only_response_finishes_hooks to use
assert_not_called() on service._gen_client.send_request instead of
assert_not_awaited(), ensuring any invocation of the generation client is
detected regardless of whether it was awaited.
Source: Path instructions
| assert result is ctx_response | ||
|
|
||
| hooks.on_resp_done.assert_called_once_with("", request, ctx_response) | ||
| service._gen_client.send_request.assert_not_awaited() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
python - <<'PY'
from unittest.mock import AsyncMock
mock = AsyncMock()
pending = mock()
mock.assert_not_awaited()
assert mock.call_count == 1
pending.close()
print("assert_not_awaited does not imply assert_not_called")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 209
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("tests/unittest/disaggregated/test_openai_disagg_service.py")
lines = path.read_text().splitlines()
for start in (360, 390, 400, 410, 430):
end = min(len(lines), start + 30)
print(f"\n--- {path}:{start+1}-{end} ---")
for i in range(start, end):
print(f"{i+1:4d}: {lines[i]}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 8079
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("tests/unittest/disaggregated/test_openai_disagg_service.py")
text = path.read_text().splitlines()
# Print the helper/implementation references around _send_disagg_request if present in same file.
for i, line in enumerate(text):
if "async def _send_disagg_request" in line or "def _send_disagg_request" in line:
start = max(0, i - 20)
end = min(len(text), i + 220)
print(f"--- {path}:{start+1}-{end} ---")
for j in range(start, end):
print(f"{j+1:4d}: {text[j]}")
break
else:
print("No _send_disagg_request implementation in this file")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 209
Use assert_not_called() here
assert_not_awaited() still passes if service._gen_client.send_request is called but never awaited. This test should enforce the no-call contract.
Proposed fix
- service._gen_client.send_request.assert_not_awaited()
+ service._gen_client.send_request.assert_not_called()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| service._gen_client.send_request.assert_not_awaited() | |
| service._gen_client.send_request.assert_not_called() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/disaggregated/test_openai_disagg_service.py` at line 407,
Replace assert_not_awaited() with assert_not_called() for
service._gen_client.send_request so the test verifies the method is never
invoked, not merely never awaited.
pcastonguay
changed the title
6 participants
Doc: TRTLLM-Serving Performance Metrics Design
Dev Engineer Review
/perf_metricspolling/endpoint behavior and moved per-request metrics delivery to:return_perf_metrics=trueandX-TRTLLM-return-metrics: 1).[DONE]metrics event (suppressed from user-visible streaming output).perf_metrics_output_dir.DisaggServerConfigaddedreturn_perf_metricsandperf_metrics_output_dir.extract_disagg_cfg(...)now propagates options intocontext_serversandgeneration_serversand forces downstream returning when eitherreturn_perf_metricsis enabled orperf_metrics_output_diris set.kwargsdirectly; switched to explicit keyword construction forDisaggServerConfig._perf_metrics_schema.pywithTypedDictrecord types supporting aggregated + disaggregated phases and a unifiedPerfMetricsRecordunion.tensorrt_llm/serve/perf_metrics.py(build_request_metrics_record,combine_disagg_metrics, header formatting, SSE event construction).openai_server.pyremoved the legacy in-memory/perf_metricsimplementation and deque/lock buffering; metrics are now built into per-request state and emitted at the correct time (explicitly awaited for both non-streaming and streaming completion).openai_disagg_server.pyreplaced prior collector behavior with queue-latency-aware hooks + conditional collection/persistence, and removed the in-process/perf_metricsroute registration.PerfMetricsMiddlewareandPerfMetricsJsonlWriter, with middleware opting-in header/SSE exposure based onX-TRTLLM-return-metrics: 1and server config.OpenAIHttpClientnow sets return-metrics headers whenrequest_perf_metricsis enabled, extracts metrics from response headers and/or the SSE post-[DONE]metrics payload, and emits via hooks.disagg_request_idvia hooks when regenerated.collect_metrics()plumbing fromOpenAIClient/OpenAIHttpClient.perf_metrics_output_dirinstead of fetching/perf_metrics.[DONE]handling in SSE consumption now exits immediately (return) to prevent extra-chunk processing.CI failures noted in context: Four reported failures—prioritize validating the new gating/emission/JSONL expectations and any removed
/perf_metricsbehaviors:#60916, pipeline#49181, commit5e0ac04#61165, pipeline#49414, commitd9ff361#61270, pipeline#49505, commitdbce2f7#61508, pipeline#49725, commitdbce2f7Recommendation for follow-up investigation: focus on test failures around (1) header/SSE gating conditions, (2) disaggregated record merging and
disagg_request_idpropagation, and (3) JSONL writer output completeness/ordering (including time-breakdown fields), since the/perf_metricsendpoint and its client/scripting usage were removed.QA Engineer Review
Test code/config changes detected (under
tests/)Integration / config updates
tests/integration/defs/accuracy/test_disaggregated_serving.pyperf_metrics_output_dir(temp-derived) instead ofperf_metrics_max_requests./perf_metricsHTTP polling helper and invocation.tests/integration/defs/disaggregated/test_disaggregated.pyperf_metrics_output_dir, validate via local JSONL artifacts, and remove prior/perf_metricsfetch/recording logic.wait_for_perf_metrics_jsonl/JSONL readers rather than the HTTP endpoint.tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yamlreturn_perf_metrics: true.tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yamlreturn_perf_metrics: true.Test utilities
tests/test_common/perf_metrics_utils.pyperf_metrics-*.jsonlfrom an output directory.read_perf_metrics_jsonl,wait_for_perf_metrics_jsonl.get_timing_metrics(...)to acceptoutput_dir.Unittest coverage
tests/unittest/disaggregated/test_disagg_utils.pytest_extract_disagg_metrics_controls.tests/unittest/llmapi/apps/_test_openai_perf_metrics.pytest_metrics_endpointtest_return_perf_metrics_and_jsonl_dumptest_streaming_metrics_require_request_opt_intests/unittest/llmapi/apps/_test_openai_chat_harmony_perf_metrics.py/perf_metricsdrain/count).kv_cache_time_output_dir.tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.pyperf_metrics_output_dirtest_return_perf_metrics_and_jsonl_dumpperf_metrics_output_dir.tests/unittest/llmapi/apps/test_request_metrics.pytests/unittest/others/test_time_breakdown.pytest_parse_disaggregated_time_breakdowntest_parse_jsonl_filetests/unittest/disaggregated/test_disagg_openai_client.pytest_internal_client_accepts_perf_metrics_header_sizetest_malformed_streaming_metrics_do_not_fail_requesttests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py: removed perf-metrics/timing-metrics plumbing.tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py: removed perf-metrics/timing-metrics plumbing.tests/unittest/llmapi/apps/test_disagg_perf_metrics_collector.py: entire module removed (collector serialization test deleted).Coverage mapping vs test lists (CI/manual)
tests/integration/test_lists/are confirmed from the available context (repo diff inspection returned no discoverable file changes in this environment)./perf_metricsto header/SSE/JSONL), so CI failures may require targeted test-list validation and/or additional waivers.Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.